@fbajraktariTT, can you review this file?

FYI @odjuricicTT, as @fbajraktariTT completed internship recently.

@jnie-TT I'm not sure that this extra logic is needed. Was a test failing without this temp fix? If so, can you provide more details?

@odjuricicTT there's an assert below that checks if the currentL1Usage is 0. This error only surfaces with my changes - it's fine without my changes because we always untilize (to_layout) before returning. However it's possible now with my changes that we will return directly without any intermediate ops between the current op and the return op, and this causes issues because we wouldn't have zeroed out currentL1Usage.

Since this function doesn't decrement l1 usage on return op, the assert will fire and say that the l1 usage is non 0. My change basically adds a check that if the consumer op is a return op, we decrement the l1 usage.

@jnie-TT Thanks! Your solution is fine for now. Just please file the followup issue and reference it in the comment.

I feel like we should have attributes on the op that denote these kind of capabilities instead of having this code be special cased for a specific op. @sdjordjevicTT thoughts?

Perhaps we should add an interface to all TTNN ops called shouldTilize that defaults to true and that ops can specialize.

Yeah that would be awesome to have, I know a lot of eltwise ops are facing a similar issue regarding data type, where some ops can typecast implicitly whereas some ops cannot. This results in the IR being misaligned with the actual runtime output.

I am thinking about this scenarios, do we have some examples?

Aren't these examples? i.e. reshape, conv2d, slice & embedding, or do you mean something else?

@sdjordjevicTT if you mean the implicit typecast ops an example would be relational binary ops vs unary ops .
Relational operations take an output_dtype that we setting to typecast implicitly within the op:

template <BinaryOpType binary_op_type>
struct RelationalBinary {
    static Tensor invoke(
        uint8_t queue_id,
        const Tensor &input_tensor_a_arg,
        const Tensor &input_tensor_b_arg,
        const std::optional<const DataType> &output_dtype = std::nullopt,
        const std::optional<MemoryConfig> &memory_config = std::nullopt,
        std::optional<Tensor> optional_output_tensor = std::nullopt,
        std::optional<unary::FusedActivations> activations = std::nullopt,
        std::optional<unary::UnaryWithParam> input_tensor_a_activation = std::nullopt);

However unary ops do not:

template <UnaryOpType... unary_op_types>
Tensor ExecuteUnary<unary_op_types...>::invoke(
    const Tensor& input_tensor,
    const std::optional<MemoryConfig>& memory_config,
    const std::optional<Tensor>& optional_output_tensor) {

And our compiler doesn't distinguish between them, i.e. for unary ops it'll still assume the output tensor of the unary op is properly typecasted to the desired data type.

As for ops that don't support implicit tilization/untilization, some examples include reshape, concat, transpose.

I believe there was a misunderstanding between us. :)

Regarding Conv, Slice, and Embedding, I'm aware that they require some inputs to be in a row-major layout. I'll address this by implementing the necessary layout workarounds. If the Metal developers decide not to support tile layout for them, then we can introduce a trait\interface to accommodate them.

Regarding the implicit conversions, I get it for the data_type, but how we are specifying whether the output is in tile\row major? By defining the optional_output_tensor? I see what can be the issue, if you have some row-major input, you want to keep it row-major output for such ops. We can think about adding the interface on an op level to support this.

I created issues on myself to follow up on this:

• [TTNN] Consider creating shouldTilize interface on TTNN op level #1862
• [TTNN] Consider creating supportImplicitTilizing interface on TTNN op level #1863

Same as above.

This will be cleaned up with the workarounds. We have tasks for each of these to cleanup.

This will be cleaned up with the workarounds. We have tasks for each of these to cleanup.

I am thinking about this scenarios, do we have some examples?

So there is strategy on LayoutDesc that will be set to ::tt::target::DistributedTensorConfig::NONE for single chip setup. Or it will be set to some kind of multi-device distribution if set. LMK if this doesn't resolve this issue.

table LayoutDesc {                                                         
  stride: [int];                                                         
  oob_val: OOBVal;
  core_range_set: [Dim2dRange];
  memory_desc: MemoryDesc;
  strategy: DistributionStrategy;
}

Reach out to @wooseokTT if you need help/interpreting its programming.

@nsmithtt the strategy doesn't tell us which submesh a tensor belongs to though right? I remember that when I added it, it was used to specify the tensor distribution method across multi device (replicate, shard etc.).

I can use it to distinguish between single/multichip, but I don't know the mesh shape or mesh offset that the tensor is mapped to if it's multidevice. And I need this info if I want to move a tensor to a multidevice mesh in the toLayout API.

It depends on the strategy, but for e.g. ShardTensor2D.shard_mesh does tell you the mesh shape. I think it's always inferred that the offset is implicitly [0, 0], @wooseokTT feel free to correct me if I'm wrong, but this is reflected from TTNN API which doesn't support arbitrary mesh offsets.

Yeah you're right ShardTensor2D has the shard_mesh. But seems like the other ones don't... If all we're using is ShardTensor2D and offset 0, 0 then I guess I can just derive it from that. And maybe add an assert that checks the strategy must be ShardTensor2D. Does doing it this way make sense with how we're performing multichip operations?

Nit: Per coding style guidelines we decided not to use alternative tokens(not, and, or, etc...):
https://docs.tenstorrent.com/tt-mlir/coding-guidelines.html#using-alternative-tokens-and-or-xor-etc

Can you please switch to regular logical symbols? I have an ongoing PR, that fixes this file:
#2055

Thanks!

Can you please open an issue for me regarding this and how I can repro it? I would like to take a look, currently, there is a canonicalization pass after the workaround pass, so I can take a look at it.

Issue created here: #2102